/****************************************************
          File Name:new_encoder.c
          Author: Abby Cin
          Mail: abbytsing@gmail.com
          Created Time: Thu 13 Nov 2014 09:36:02 PM CST
*****************************************************/

#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <stdio.h>
#include <fcntl.h>
#include <unistd.h>
#include <x264.h>

void trans_from(int width, int height, uint8_t *yuv422_frame, uint8_t *yuv420_frame)
{
	int i, j, k =0;
	int size = width * height;
	/* get Y */
	for(i = 0; i < size; i++)
		yuv420_frame[i] = yuv422_frame[i*2];
	
	/* get U */
	for(i = 0; i < height; i++)
	{
		if((i % 2) != 0)
			continue;
		for(j = 0; j < (width / 2); j++)
		{
			if((4 * j + 1) > (2 * width))
				break;
			yuv420_frame[size + k * 2 * width / 4 + j] = yuv422_frame[i * 2 * width + 4 * j +1];
		}
		k++;
	}
	
	k = 0;
	
	/* get V */
	for(i = 0; i < height; i++)
	{
		if((i % 2) == 0)
			continue;
		for(j = 0; j < (width / 2); j++)
		{
			if((4 * j + 3) > (2 * width))
				break;
			yuv420_frame[size + size / 4 + k * 2 * width / 4 + j] = yuv422_frame[i * 2 * width + 4 * j + 3];
		}
		k++;
	}
}

void compress_frame(x264_picture_t *pic_in, uint8_t *yuv420_frame, int width, int height)
{
	pic_in->img.plane[0] = yuv420_frame;
	pic_in->img.plane[1] = pic_in->img.plane[0] + width * height;
	pic_in->img.plane[2] = pic_in->img.plane[1] + width * height / 4;
}

int main(int argc, char *argv[])
{
	if(argc < 3 || strcmp(argv[1],"--help") == 0)
	{
		fprintf(stderr,"%s in.yuv out.h264\n",argv[0]);
		exit(1);
	}
	
	int width = 640;
	int height = 480;
	int j = 1;									/* count from 1 */
	int inf, outf;
	int nnal = -1; 									/* segments in pic_in */
	x264_nal_t *nals, *nal;
	x264_t *encoder;
	x264_picture_t *pic_in = NULL, pic_out;
	x264_param_t *param;
	
	size_t yuv422_size = width * height * 2;			/* yuv422 length */
	size_t yuv420_size = width * height * 3 / 2;		/* yuv420 length */
    uint8_t *yuv422_frame = (uint8_t *) malloc(yuv422_size);
	uint8_t *yuv420_frame = (uint8_t *) malloc(yuv420_size);
	
/* initailize fd */
	inf = open(argv[1], O_RDONLY);
	if (inf < 0) {
		return -1;
	}
	outf = open(argv[2], O_CREAT | O_WRONLY | O_APPEND, 00644);
	if (outf < 0) {
		return -1;
	}
	
/* initailize encoder */
	param = (x264_param_t *) malloc(sizeof(x264_param_t));
	
	x264_param_default(param); 
	x264_param_default_preset(param, "veryfast", "zerolatency");
	
	param->i_threads = 4;
	param->i_width = width;
	param->i_height = height;
	param->i_fps_num = 5;
	param->i_fps_den = 1;
	param->rc.i_lookahead = 0; 	/* buffering ahead */
	param->i_keyint_max = 300;
	param->i_keyint_min = X264_KEYINT_MIN_AUTO;

	x264_param_apply_profile(param, "baseline");
	
	encoder = x264_encoder_open(param);
	
/* initailize pic_in */
	pic_in = (x264_picture_t *) malloc(sizeof(x264_picture_t));
	x264_picture_alloc(pic_in, X264_CSP_I420, width, height);
	pic_in->img.i_csp = X264_CSP_I420;
	pic_in->img.i_plane = 3;
	
/* encoding */
	while (read(inf, yuv422_frame, yuv422_size) > 0) 
	{
		trans_from(width, height, yuv422_frame, yuv420_frame);
		
		compress_frame(pic_in, yuv420_frame, width, height);
		
		if(x264_encoder_encode(encoder, &nals, &nnal, pic_in, &pic_out))
		{
			for(nal = nals; nal < nals + nnal; nal++)
				write(outf, nal->p_payload, nal->i_payload);
			printf("\e[32m --------- %dth frame --------\e[0m\n",j++);
		}
		else
		{
			perror("encoding failed");
			exit(errno);
		}
	}
	
/* clean */
	x264_encoder_close(encoder);
	close(inf);
	close(outf);
	free(yuv422_frame);
	free(yuv420_frame);
	return 0;
}
